{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cfb1f22c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/lib/python3.10/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: '/root/miniconda3/lib/python3.10/site-packages/torchvision/image.so: undefined symbol: _ZN3c1017RegisterOperatorsD1Ev'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?\n",
      "  warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading pretrained fmde from checkpoints/depth_anything_v2_vitl.pth...\n",
      "Freezing the decoder now.\n",
      "Freezing the encoder now.\n",
      "Loading pretrained cmde from checkpoints/depth_anything_v2_vitb.pth...\n",
      "Loading checkpoint from checkpoints/prior_depth_anything_vitb.pth...\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import os \n",
    "\n",
    "from os.path import join, exists\n",
    "\n",
    "\n",
    "import cv2\n",
    "\n",
    "from PIL import Image\n",
    "\n",
    "from daniel_tools.img_utils import * \n",
    "\n",
    "from prior_depth_anything.utils import colorize_depth_maps,chw2hwc\n",
    "\n",
    "    \n",
    "\n",
    "import torch\n",
    "from prior_depth_anything import PriorDepthAnything\n",
    "\n",
    "\n",
    "def depth2colordepth(gt_depth):\n",
    "    value_colored = colorize_depth_maps((gt_depth  - gt_depth.min()) / (gt_depth.max() - gt_depth.min()), 0, 1, cmap=\"Spectral\").squeeze()\n",
    "    value_colored = (value_colored * 255).astype(np.uint8)\n",
    "    value_colored = Image.fromarray(chw2hwc(value_colored))\n",
    "\n",
    "    return value_colored\n",
    "\n",
    "\n",
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "ckpt_dir ='checkpoints'\n",
    "\n",
    "\n",
    "# priorda = PriorDepthAnything(device=device, fmde_dir=ckpt_dir, ckpt_dir=ckpt_dir,cmde_dir = ckpt_dir)\n",
    "priorda = PriorDepthAnything(device=device, fmde_dir=ckpt_dir, ckpt_dir=ckpt_dir,cmde_dir = ckpt_dir, frozen_model_size = 'vitl', conditioned_model_size='vitb')\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4571e9f9",
   "metadata": {},
   "source": [
    "# load depth-pro-prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90e10ec2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "\n",
    "class InferDepthLoader(Dataset):\n",
    "\n",
    "    def __init__(self, jsons=['test.jsonl']):\n",
    "        super().__init__()\n",
    "        \n",
    "        all_data = []\n",
    "        for json_name in jsons:\n",
    "            root = dirname(json_name)\n",
    "            with open(os.path.join(root,json_name ), \"r\") as f:\n",
    "                for line in f:\n",
    "                    data_term = json.loads(line)\n",
    "                    # data_term['image'] = os.path.join(root, data_term['image'])\n",
    "                    if data_term.get('conditioning_image',None) is not None :\n",
    "                        data_term['conditioning_image'] = os.path.join(root, data_term['conditioning_image'])\n",
    "                    #!  for tricky\n",
    "                    else:\n",
    "                        data_term['conditioning_image'] = os.path.join(root, data_term['data_path'])\n",
    "\n",
    "                    all_data += [data_term]\n",
    "                \n",
    "\n",
    "        \n",
    "        logger.warning(f\"data size is {len(all_data)}\")\n",
    "        \n",
    "\n",
    "        self.all_data = all_data\n",
    "\n",
    "\n",
    "        self.transforms = None\n",
    "        # self.transforms = torchvision.transforms.Compose([\n",
    "        #     torchvision.transforms.ToTensor(),\n",
    "        # ])\n",
    "\n",
    "    def __len__(self, ):\n",
    "        return len(self.all_data)\n",
    "\n",
    "    def __getitem__(self,idx):\n",
    "\n",
    "        item = self.all_data[idx]\n",
    "\n",
    "        input_image_path = item['conditioning_image']\n",
    "        \n",
    "\n",
    "        input_image_dir = input_image_path.split('/')[-2]\n",
    "        input_image_name = input_image_path.split('/')[-1][:-4]\n",
    "        \n",
    "        \n",
    "        test_image = Image.open(input_image_path).convert('RGB')\n",
    "        test_image = np.array(test_image).astype(np.float32)\n",
    "        test_image = test_image / 127.5 - 1.0 \n",
    "        if self.transforms is not None:\n",
    "            #* B, H, W\n",
    "            test_image = self.transforms(test_image) #* B, H, W\n",
    "        \n",
    "\n",
    "\n",
    "        #* depth uint is m \n",
    "        return dict(hint = test_image,hint_path =item['conditioning_image'], )\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e090d70b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m2025-06-26 22:00:14.330\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m25\u001b[0m - \u001b[33m\u001b[1mdata size is 369\u001b[0m\n",
      "  0%|          | 0/369 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/share/project/cwm/shaocong.xu/exp/Prior-Depth-Anything/prior_depth_anything/depth_completion.py:153: UserWarning: The depth prior is directly provided by the user. All the known points will cover the knn-scaled map.\n",
      "  warnings.warn(\"The depth prior is directly provided by the user. All the known points will cover the knn-scaled map.\")\n",
      "100%|██████████| 369/369 [00:50<00:00,  7.33it/s]\n"
     ]
    }
   ],
   "source": [
    "from os.path import dirname\n",
    "import json\n",
    "from loguru import logger\n",
    "import torchvision\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "\n",
    "prior_root = '/share/project/cwm/shaocong.xu/exp/Prior-Depth-Anything/data/tricky_test_moge'\n",
    "\n",
    "\n",
    "loader = InferDepthLoader(jsons=['/share/project/cwm/shaocong.xu/exp/Prior-Depth-Anything/data/tricky_testset/test.jsonl'])\n",
    "\n",
    "\n",
    "target_root = '/share/project/cwm/shaocong.xu/exp/Prior-Depth-Anything/logs/PriorDA_MoGe_test_mogev2'\n",
    "\n",
    "\n",
    "\n",
    "os.makedirs(target_root, exist_ok=True)\n",
    "\n",
    "\n",
    "\n",
    "for batch in tqdm(loader):\n",
    "    # print(batch['hint'].shape, batch['hint_path'])\n",
    "    tmp = batch['hint_path'].split('/')\n",
    "\n",
    "    # path = join(prior_root,tmp[-3],tmp[-1].replace('png','npz'))\n",
    "    # tgt_path = join(prior_root,tmp[-3],tmp[-1].replace('png','npy'))\n",
    "    # np.save(tgt_path,np.load(path)['depth'].astype('float32'))\n",
    "\n",
    "    prior_depth_path = join(prior_root,tmp[-3],tmp[-1].replace('png','npy'))\n",
    "\n",
    "    save_root = join(target_root,tmp[-3])\n",
    "    os.makedirs(save_root, exist_ok=True)\n",
    "    save_path = join(save_root, tmp[-1].replace('png','npy'))\n",
    "\n",
    "    # print(f'save_path: {save_path}')\n",
    "    \n",
    "    # print(batch['hint_path'], '\\n', prior_depth_path,)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    # output = priorda.infer_one_sample(image=batch['hint_path'], prior=prior_depth_path,visualize=True)\n",
    "    output = priorda.infer_one_sample(image=batch['hint_path'], prior=prior_depth_path)\n",
    "\n",
    "\n",
    "\n",
    "    # depth2colordepth(output.cpu().numpy())\n",
    "\n",
    "    np.save(save_path, output.cpu().numpy().astype('float32'))\n",
    "\n",
    "\n",
    "    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3dc4196",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(752, 1028)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5116a82d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
